import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPRegressor
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.preprocessing import StandardScaler
from common_import import *


def add_bias_node(X):
    # 创建一个与X行数相同的列，值为1
    bias_column = np.ones((X.shape[0], 1))
    # 将偏倚列添加到X的前面
    X_with_bias = np.hstack((bias_column, X))
    return X_with_bias


def plot_predictions(y_test, y_pred):
    plt.figure(figsize=(10, 6))

    # 绘制预测值与实际值
    plt.scatter(y_test, y_pred, label="预测值", color="blue", alpha=0.6)

    # 绘制 y_test = y_pred 的参考线
    plt.plot(y_test, y_test, color="red", linestyle="--", label="y_test = y_pred")

    # 设置坐标轴刻度相同
    plt.xlim(min(y_test), max(y_test))
    plt.ylim(min(y_test), max(y_test))

    # 添加标签和标题
    plt.xlabel("实际值")
    plt.ylabel("预测值")
    plt.title("")
    plt.legend()
    plt.grid()

    # 显示图形
    plt.show()


def run_mlp_regressor(X, y):
    """
    使用多层感知器回归模型来预测目标值
    """

    X_with_bias = add_bias_node(X)

    # 数据标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X_with_bias)

    # 划分数据集为训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X_scaled, y, test_size=0.2, random_state=42
    )

    # 定义MLP回归模型
    mlp = MLPRegressor(
        hidden_layer_sizes=(200, 80),  # 隐藏层结构，可以根据需要调整
        activation="relu",  # 激活函数，默认ReLU
        solver="adam",  # 优化器，使用Adam
        max_iter=500,  # 最大迭代次数，可以调整以获得更好的收敛
        random_state=42,
    )

    # 训练模型
    mlp.fit(X_train, y_train)

    # 对测试集进行预测
    y_pred = mlp.predict(X_test)

    # 模型评估
    mse = mean_squared_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    print(f"Mean Squared Error (MSE): {mse:.4f}")
    print(f"R^2 Score: {r2:.4f}")
    y_pred1 = mlp.predict(X_train)
    amse = mean_squared_error(y_train, y_pred1)
    ar2 = r2_score(y_train, y_pred1)
    print(f"Mean Squared Error (MSE): {amse:.4f}")
    print(f"R^2 Score: {ar2:.4f}")


def predict_with_mlp(X, y):
    # 数据标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # 划分数据集为训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X_scaled, y, test_size=0.2, random_state=42
    )

    # 定义MLP回归模型
    mlp = MLPRegressor(
        hidden_layer_sizes=(100, 50),  # 隐藏层结构
        activation="relu",  # 激活函数
        solver="adam",  # 优化器
        max_iter=500,  # 最大迭代次数
        random_state=42,
    )

    # 训练模型
    mlp.fit(X_train, y_train)

    # 对测试集进行预测
    y_pred = mlp.predict(X_test)

    # 模型评估
    mse = mean_squared_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    print("均方误差 (MSE):", mse)
    print("决定系数 (R²):", r2)
    tool.plot_predictions(y_test, y_pred, "mlp预测.png")
    return y_pred


if __name__ == "__main__":
    molecular_descriptor = pd.read_csv("data/Molecular_Descriptor_training.csv")
    era_activity = pd.read_csv("data/ER_activity_training.csv")

    # 移除不必要的列（例如化合物ID）
    X = molecular_descriptor.drop(columns=["SMILES"])
    y = era_activity["pIC50"]

    filter_X = X[constants.feature_20]

    # 运行MLP回归
    predict_with_mlp(filter_X, y)
